package com.fanshuai.soring.boot.webflux;

import com.fanshuai.soring.boot.contants.CommonConstants;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.reactive.result.method.RequestMappingInfo;
import org.springframework.web.reactive.result.method.annotation.RequestMappingHandlerMapping;
import org.springframework.web.server.ServerWebExchange;

import java.util.List;
import java.util.Map;
import java.util.Set;

@Slf4j
public class MyHandlerMapping extends RequestMappingHandlerMapping implements InitializingBean {
    private HandlerMethod notFoundMethod = null;

    @Override
    public void afterPropertiesSet() {
        super.afterPropertiesSet();

        Map<RequestMappingInfo, HandlerMethod> handlerMethodMap = super.getHandlerMethods();
        for (RequestMappingInfo info : handlerMethodMap.keySet()) {
            if (info.getPatternsCondition().getPatterns().stream().anyMatch(pathPattern -> {
                return pathPattern.getPatternString().equals(CommonConstants.notFound);
            })) {
                notFoundMethod = handlerMethodMap.get(info);
                log.info("set notFoundMethod={}", notFoundMethod);
            }
        }
    }

    @Override
    protected HandlerMethod handleNoMatch(Set<RequestMappingInfo> infos, ServerWebExchange exchange) throws Exception {
        String uri = exchange.getRequest().getURI().getPath();
        String metricsHeader = getHeader("metrics", exchange);
        boolean metrics = uri.startsWith("/metrics") || StringUtils.isNotBlank(metricsHeader);

        if (!metrics) {
            exchange.getAttributes().put(CommonConstants.ORIGINAL_URI, exchange.getRequest().getURI());
            return notFoundMethod;
        }
        return super.handleNoMatch(infos, exchange);
    }

    private String getHeader(String headerName, ServerWebExchange exchange) {
        List<String> metricsHeader = exchange.getRequest().getHeaders().get("metrics");
        return CollectionUtils.isEmpty(metricsHeader) ? null : metricsHeader.get(0);
    }
}
